package data

import (
	"encoding/json"
	"fmt"
	"log"
	"sync"

	"github.com/Shopify/sarama"
)

// KafkaConsumer 处理Kafka消息消费
type KafkaConsumer struct {
	consumer   sarama.Consumer
	partitions map[string][]int32
	wg         sync.WaitGroup
	stopChan   chan struct{}
}

// NewKafkaConsumer 创建一个新的Kafka消费者
func NewKafkaConsumer(kafkaBrokers []string) (*KafkaConsumer, error) {
	consumer, err := sarama.NewConsumer(kafkaBrokers, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create kafka consumer: %v", err)
	}

	return &KafkaConsumer{
		consumer:   consumer,
		partitions: make(map[string][]int32),
		stopChan:   make(chan struct{}),
	}, nil
}

// ConsumeTopic 开始消费指定的主题
func (kc *KafkaConsumer) ConsumeTopic(topic string, handler func([]byte) error) error {
	// 获取主题的所有分区
	partitions, err := kc.consumer.Partitions(topic)
	if err != nil {
		return fmt.Errorf("failed to get partitions for topic %s: %v", topic, err)
	}

	kc.partitions[topic] = partitions

	// 为每个分区创建一个消费者
	for _, partition := range partitions {
		kc.wg.Add(1)
		go func(p int32) {
			defer kc.wg.Done()
			kc.consumePartition(topic, p, handler)
		}(partition)
	}

	log.Printf("Started consuming topic %s with %d partitions", topic, len(partitions))
	return nil
}

// consumePartition 消费指定分区的消息
func (kc *KafkaConsumer) consumePartition(topic string, partition int32, handler func([]byte) error) {
	// 从分区的最新偏移量开始消费
	partitionConsumer, err := kc.consumer.ConsumePartition(topic, partition, sarama.OffsetNewest)
	if err != nil {
		log.Printf("Failed to start consumer for partition %d: %v", partition, err)
		return
	}
	defer partitionConsumer.Close()

	log.Printf("Started consuming partition %d of topic %s", partition, topic)

	for {
		select {
		case msg := <-partitionConsumer.Messages():
			// 调用消息处理函数
			err := handler(msg.Value)
			if err != nil {
				log.Printf("Error processing message from partition %d: %v", partition, err)
			}
		case <-kc.stopChan:
			log.Printf("Stopping consumer for partition %d of topic %s", partition, topic)
			return
		}
	}
}

// Stop 停止所有消费者
func (kc *KafkaConsumer) Stop() {
	// 发送停止信号
	close(kc.stopChan)

	// 等待所有消费者goroutine结束
	kc.wg.Wait()

	// 关闭消费者
	err := kc.consumer.Close()
	if err != nil {
		log.Printf("Warning: failed to close kafka consumer: %v", err)
	}

	log.Println("All kafka consumers stopped")
}

// ModelTrainingHandler 处理用于模型训练的数据
func ModelTrainingHandler(data []byte) error {
	// 这里可以实现将数据传递给TensorFlow进行模型训练的逻辑
	// 由于我们只是创建一个框架，这里仅打印日志
	var msgData map[string]interface{}
	if err := json.Unmarshal(data, &msgData); err != nil {
		return fmt.Errorf("failed to unmarshal message data: %v", err)
	}

	// 模拟将数据发送给TensorFlow进行模型训练
	deviceID := msgData["device_id"].(string)
	timestamp := msgData["timestamp"].(float64)

	log.Printf("Sending data to TensorFlow for training: device=%s, timestamp=%v", deviceID, timestamp)

	// 实际应用中，这里应该调用TensorFlow的API或通过其他方式将数据传递给TensorFlow

	return nil
}
